"""Moonshot assistant"""

from dataclasses import dataclass, asdict
from typing import List

from openai import OpenAI
from loguru import logger


@dataclass
class Message:
    role: str
    content: str


class MoonshotAssistant:
    def __init__(self, api_key, base_url="https://api.moonshot.cn/v1"):
        self.client = OpenAI(api_key=api_key, base_url=base_url)

    def chat(
        self, messages: List[Message], temperature: float = 0.3, debug: bool = False
    ):
        messages = [asdict(m) for m in messages]
        completion = self.client.chat.completions.create(
            model="moonshot-v1-8k",
            messages=messages,
            temperature=temperature,
        )
        ret = completion.choices[0].message.content
        if debug:
            logger.debug(ret)
        return ret
